import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori,association_rules
import matplotlib.pyplot as plt
import networkx as nx

df = pd.read_excel('associationRules/餐厅数据.xlsx')

#print (df.head)
#转换菜品数据格式
trsations = df['菜品'].str.split(',').to_list()

#标准化数据
ts = TransactionEncoder()
te_ary = ts.fit(trsations).transform(trsations)
df_enecoded = pd.DataFrame(te_ary,columns=ts.columns_)

#print(df_enecoded.head)
#使用apriori进行关联规则挖掘
frequent_itemsets = apriori(df_enecoded,min_support=0.1,use_colnames=True)
frequent_itemsets.sort_values(by='surport',ascending=False,inplace=True)

#选择2项集查看
# print(frequent_itemsets[frequents_itemsets.itemsets.apply(lambda x : len(x) == 2)])

rules = association_rules(frequent_itemsets, metric='confidence',min_threshold=0.1)
#筛选有效规则
effective=rules[
    (rules['lift'] > 1) & (rules['confidence'] > 0.1)
].sort_values(by=['lift', 'confidence'],ascending=False)

print(effective[['antecedents', 'consequents', 'support', 'confidence','lift']])

#设置字体
plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus']=False
#关联关系可视化
G = nx.DiGraph()
for _, row in effective.iterrows():
    G.add_edge(','.join(list(row['antecedents'])),
               ','.join(list(row['consequents'])),
               weight=row['lift']
               )
plt.figure(figsize=(12,8))
pos=nx.spring_layout(G)
nx.draw(G,pos,
        with_labels=True,
        edge_color=[G[u][v]['weight'] for u, v in G.edges()],
        width=2.0,
        edge_cmap=plt.cm.Blues)
plt.show()